# _*_ coding:utf-8 _*_
"""
@author:lx
@file: segement3
@contact: lixiang-929@outlook.com
@time: 2019/03/23
"""
import cv2
import numpy as np
import os
import time

data_path = ".//data//"
result_path = ".//result//"

def segement(data_path, result_path):
    data_list = os.listdir(data_path)
    for data in data_list:
        start_time = time.time()
        image_name = os.path.join(data_path, data)
        image = cv2.imread(image_name)
        image_width = image.shape[1]
        image_height = image.shape[0]
        b_channel, g_channel, r_channel = cv2.split(image)

        rg_image = np.zeros([image_height, image_width])
        rb_image = np.zeros([image_height, image_width])
        gb_image = np.zeros([image_height, image_width])
        gr_image = np.zeros([image_height, image_width])
        fR = np.zeros([image_height, image_width])
        fY = np.zeros([image_height, image_width])
        fXY = np.zeros([image_height, image_width])
        for j in range(image_height):
            for i in range(image_width):
                rg_image[j, i] = int(r_channel[j, i]) - int(g_channel[j, i])
                rb_image[j, i] = int(r_channel[j, i]) - int(b_channel[j, i])
                # gr_image[j, i] = int(g_channel[j, i]) - int(r_channel[j, i])
                gb_image[j, i] = int(g_channel[j, i]) - int(b_channel[j, i])

                fR[j, i] = max(0, min(rg_image[j, i], rb_image[j, i]))
                fY[j, i] = max(0, min(rb_image[j, i], gb_image[j, i]))
                fXY[j, i] = fR[j, i] + fY[j, i]

        file_name, file_format = os.path.splitext(data)
        fR_name = os.path.join(result_path, file_name + "_fR.jpg")
        fY_name = os.path.join(result_path, file_name + "_fY.jpg")
        fXY_name = os.path.join(result_path,file_name + "_fXY.jpg")
        cv2.imwrite(fR_name, fR)
        cv2.imwrite(fY_name, fY)
        cv2.imwrite(fXY_name,fXY)
        print("{0:.3f}, {1:s}".format((time.time() - start_time), file_name))
        cv2.waitKey(10)

if __name__ == "__main__":
    segement(data_path,result_path)